#include "codec_ctx.h"
#include <assert.h>
#include <utility>
#include <functional>

codec_context::codec_context()
{
}

codec_context::~codec_context()
{
  if (parser_)
  {
    av_parser_close(parser_);
  }
  if (c_)
  {
    avcodec_free_context(&c_);
  }
  if (frame_)
  {
    av_frame_free(&frame_);
  }

  if (pkt_)
  {
    av_packet_free(&pkt_);
  }

  if (render_frame_)
  {
    av_frame_free(&render_frame_);
  }
  Destroy_Ring_Queue(iq);
  Destroy_Ring_Queue(oq);
  if (decoder_thread_.joinable())
  {
    quit_ = true;
    {
      std::unique_lock<std::mutex> lg(cv_mu_);
      cv_.notify_all();
    }
    decoder_thread_.join();
  }
}

bool codec_context::create_sws_ctx(int width, int height, AVPixelFormat src_fmt)
{
  sws_ctx_ = sws_getContext(width, height, src_fmt, width, height, AV_PIX_FMT_RGB24, SWS_BICUBIC, NULL, NULL, NULL);
  if (!sws_ctx_)
  {
    return false;
  }
  render_frame_ = av_frame_alloc();
  int dst_bts_num = avpicture_get_size(AV_PIX_FMT_RGB24, width, height);
  render_buffer_.resize(dst_bts_num);
  avpicture_fill((AVPicture *)render_frame_, render_buffer_.data(), AV_PIX_FMT_RGB24, width, height);
  return true;
}

bool codec_context::create(bool enable_asyn)
{
  pkt_ = av_packet_alloc();
  if (!pkt_)
  {
    return false;
  }
  codec_ = avcodec_find_decoder(AV_CODEC_ID_HEVC);
  if (!codec_)
  {
    return false;
  }

  parser_ = av_parser_init(codec_->id);
  if (!parser_)
  {
    return false;
  }
  c_ = avcodec_alloc_context3(codec_);
  if (!c_)
  {
    return false;
  }
  if (avcodec_open2(c_, codec_, NULL) < 0)
  {
    return false;
  }
  frame_ = av_frame_alloc();
  if (!frame_)
  {
    return false;
  }
  // sem count
  iq = Create_Ring_Queue(20, sizeof(void *));
  oq = Create_Ring_Queue(2, sizeof(void *));
  current_frame_cnt_.store(0, std::memory_order::memory_order_relaxed);
  if (enable_asyn)
  {
    decoder_thread_ = std::thread(std::bind(&codec_context::decode_thread_loop, this));
  }
  valid_ = true;
  return true;
}

void codec_context::decode_thread_loop()
{
  printf("Decoder thread start!\n");
  while (!quit_)
  {
    if (current_frame_cnt_.load(std::memory_order::memory_order_acquire) == 0)
    {
      std::unique_lock<std::mutex> lg(cv_mu_);
      cv_.wait(lg, [this]()
               { return (current_frame_cnt_.load(std::memory_order::memory_order_acquire) > 0) || quit_; });
      if (quit_)
      {
        break;
      }
    }
    current_frame_cnt_.fetch_sub(1, std::memory_order::memory_order_acq_rel);
    std::vector<uint8_t> **ppvec = (std::vector<uint8_t> **)SOLO_Read(iq);
    assert(ppvec);
    auto pvec = *ppvec;
    *ppvec = NULL;
    SOLO_Read_Over(iq);

    std::vector<std::vector<uint8_t>> *res = new std::vector<std::vector<uint8_t>>();
    decode_frames(pvec->data(), pvec->size(), *res);
    delete pvec;
    if (res->empty())
    {
      delete res;
      continue;
    }
    std::vector<std::vector<uint8_t>> **ppovec = (std::vector<std::vector<uint8_t>> **)SOLO_Write(oq);
    if (!ppovec)
    {
      printf("Decoder drop!\n");
      delete res;
      continue;
    }
    *ppovec = res;
    SOLO_Write_Over(oq);
  }
  printf("Decoder thread stop!\n");
}

enum NALUnitType
{
  NAL_TRAIL_N = 0,
  NAL_TRAIL_R = 1,
  NAL_TSA_N = 2,
  NAL_TSA_R = 3,
  NAL_STSA_N = 4,
  NAL_STSA_R = 5,
  NAL_RADL_N = 6,
  NAL_RADL_R = 7,
  NAL_RASL_N = 8,
  NAL_RASL_R = 9,
  NAL_BLA_W_LP = 16,
  NAL_BLA_W_RADL = 17,
  NAL_BLA_N_LP = 18,
  NAL_IDR_W_RADL = 19,
  NAL_IDR_N_LP = 20,
  NAL_CRA_NUT = 21,
  NAL_VPS = 32,
  NAL_SPS = 33,
  NAL_PPS = 34,
  NAL_AUD = 35,
  NAL_EOS_NUT = 36,
  NAL_EOB_NUT = 37,
  NAL_FD_NUT = 38,
  NAL_SEI_PREFIX = 39,
  NAL_SEI_SUFFIX = 40,
};

static int parse_265_nalu(char *data, int len)
{
  unsigned int code = -1;
  int vps = 0, sps = 0, pps = 0, irap = 0;
  int i;
  for (i = 0; i < len - 1; i++)
  {
    code = (code << 8) + data[i];
    if ((code & 0xffffff00) == 0x100)
    {
      char nal2 = data[i + 1];
      int type = (code & 0x7E) >> 1;

      if (code & 0x81) // forbidden and reserved zero bits
        return 0;

      if (nal2 & 0xf8) // reserved zero
        return 0;

      switch (type)
      {
      case NAL_VPS:
        vps++;
        break;
      case NAL_SPS:
        sps++;
        break;
      case NAL_PPS:
        pps++;
        break;
      case NAL_BLA_N_LP:
      case NAL_BLA_W_LP:
      case NAL_BLA_W_RADL:
      case NAL_CRA_NUT:
      case NAL_IDR_N_LP:
      case NAL_IDR_W_RADL:
        irap++;
        break;
      }
    }
  }
  return 0;
}

int codec_context::decode_frames(uint8_t *data, int len, std::vector<std::vector<uint8_t>> &res)
{
  int ret;
  res.clear();
  while (len > 0)
  {
    int ret = av_parser_parse2(parser_, c_, &pkt_->data, &pkt_->size, data, len, AV_NOPTS_VALUE, AV_NOPTS_VALUE, 0);
    if (ret < 0)
    {
      return -1;
    }
    len -= ret;
    data += ret;
    // ʹ���˶�������
    // used = ret;
    if (!pkt_->size)
    {
      continue;
    }
    ret = avcodec_send_packet(c_, pkt_);
    if (ret < 0)
    {
      fprintf(stderr, "Error sending a packet for decoding\n");
      return -1;
    }
    for (;;)
    {
      ret = avcodec_receive_frame(c_, frame_);
      if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
      {
        break;
      }
      else if (ret < 0)
      {
        fprintf(stderr, "Error during decoding\n");
        return -1;
      }
      if (!sws_ctx_)
      {
        create_sws_ctx(frame_->width, frame_->height, c_->pix_fmt);
      }
      convert();
      // dec_cb_(render_buffer_.data(), render_buffer_.size(), render_frame_->linesize[0]);
      std::vector<uint8_t> result(frame_->width * 3 * frame_->height);
      int step = frame_->width * 3;
      for (int i = 0; i < frame_->height; ++i)
      {
        memcpy(result.data() + i * step, render_buffer_.data() + i * render_frame_->linesize[0], step);
      }
      res.push_back(result);
    }
  }
  return 0;
}

int codec_context::commit_packet(uint8_t *data, int len)
{
  std::vector<uint8_t> **pp_input = (std::vector<uint8_t> **)SOLO_Write(iq);
  if (!pp_input)
  {
    printf("Can't put frame data!\n");
    return -1;
  }
  std::vector<uint8_t> *p_input = new std::vector<uint8_t>(len);
  memcpy(p_input->data(), data, len);
  *pp_input = p_input;
  SOLO_Write_Over(iq);
  current_frame_cnt_.fetch_add(1, std::memory_order::memory_order_acq_rel);
  {
    std::unique_lock<std::mutex> lg(cv_mu_);
    cv_.notify_one();
  }
  return 0;
}

std::vector<uint8_t> codec_context::get_last_decoded_frame()
{
  std::vector<std::vector<uint8_t>> **pp_output = (std::vector<std::vector<uint8_t>> **)SOLO_Read(oq);
  if (!pp_output)
  {
    printf("No frame avaliable!\n");
    return {};
  }
  std::vector<std::vector<uint8_t>> *p_out = *pp_output;
  *pp_output = NULL;
  SOLO_Read_Over(oq);
  assert(!p_out->empty());
  std::vector<uint8_t> res;
  res.swap(p_out->back());
  delete p_out;
  return std::move(res);
}